"""
19_populate_ontology.py
=======================
Populate Knowledge Graph from Ontology Data

This module covers:
1. Loading ontology data from various sources
2. Batch population of Neo4j
3. Creating relationships from data
4. Validation and verification
5. Performance optimization for large datasets

Author: Neo4j Tutorial
Date: January 2025
"""

import json
import csv
from typing import Dict, List, Any, Optional, Generator
from dataclasses import dataclass
from pathlib import Path
from neo4j import Driver
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


class OntologyPopulator:
    """
    Populates Neo4j knowledge graph from ontology data sources.
    
    Data Flow:
    
    ┌──────────────────────────────────────────────────────────────────┐
    │                    DATA POPULATION PIPELINE                      │
    ├──────────────────────────────────────────────────────────────────┤
    │                                                                  │
    │  ┌─────────────┐    ┌─────────────┐    ┌─────────────┐          │
    │  │  Data       │    │  Transform  │    │   Load      │          │
    │  │  Sources    │───►│  & Validate │───►│   to Neo4j  │          │
    │  └─────────────┘    └─────────────┘    └─────────────┘          │
    │        │                   │                   │                 │
    │        ▼                   ▼                   ▼                 │
    │  - JSON files         - Type conversion   - Batch UNWIND        │
    │  - CSV files          - Null handling     - Transaction mgmt    │
    │  - API responses      - Validation        - Index usage         │
    │  - Database dumps     - Normalization     - Constraint checks   │
    │                                                                  │
    └──────────────────────────────────────────────────────────────────┘
    """
    
    def __init__(self, driver: Driver, batch_size: int = 1000):
        self.driver = driver
        self.batch_size = batch_size
        self.stats = {
            "nodes_created": 0,
            "relationships_created": 0,
            "errors": []
        }
    
    # =========================================================
    # DATA LOADING
    # =========================================================
    
    def load_json_file(self, filepath: str) -> Dict[str, Any]:
        """Load data from JSON file."""
        with open(filepath, 'r', encoding='utf-8') as f:
            return json.load(f)
    
    def load_csv_file(self, filepath: str) -> Generator[Dict, None, None]:
        """Load data from CSV file as generator."""
        with open(filepath, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            for row in reader:
                yield row
    
    def batch_generator(
        self,
        items: List[Any],
        batch_size: int = None
    ) -> Generator[List[Any], None, None]:
        """Split items into batches."""
        size = batch_size or self.batch_size
        for i in range(0, len(items), size):
            yield items[i:i + size]
    
    # =========================================================
    # POPULATE PATIENTS
    # =========================================================
    
    def populate_patients(self, patients: List[Dict]) -> int:
        """
        Populate Patient nodes from data.
        
        Args:
            patients: List of patient dictionaries
            
        Returns:
            Number of patients created
        """
        query = """
        UNWIND $batch AS p
        MERGE (patient:Patient {id: p.id})
        SET patient.name = p.name,
            patient.dateOfBirth = date(p.dateOfBirth),
            patient.gender = p.gender,
            patient.bloodType = p.bloodType,
            patient.allergies = p.allergies,
            patient.active = coalesce(p.active, true),
            patient.createdAt = datetime(),
            patient.source = 'ontology_import'
        RETURN count(*) as created
        """
        
        total = 0
        with self.driver.session() as session:
            for batch in self.batch_generator(patients):
                result = session.run(query, batch=batch)
                total += result.single()["created"]
        
        self.stats["nodes_created"] += total
        return total
    
    # =========================================================
    # POPULATE DISEASES (ICD-11)
    # =========================================================
    
    def populate_diseases(self, diseases: List[Dict]) -> int:
        """
        Populate Disease nodes from ICD-11 data.
        
        Args:
            diseases: List of disease dictionaries with ICD-11 codes
            
        Returns:
            Number of diseases created
        """
        query = """
        UNWIND $batch AS d
        MERGE (disease:Disease:ICD11Disease {id: d.id})
        SET disease.name = d.name,
            disease.category = d.category,
            disease.description = d.description,
            disease.severity = d.severity,
            disease.ontology = 'ICD-11',
            disease.createdAt = datetime()
        RETURN count(*) as created
        """
        
        total = 0
        with self.driver.session() as session:
            for batch in self.batch_generator(diseases):
                result = session.run(query, batch=batch)
                total += result.single()["created"]
        
        self.stats["nodes_created"] += total
        return total
    
    # =========================================================
    # POPULATE MEDICATIONS (RxNorm)
    # =========================================================
    
    def populate_medications(self, medications: List[Dict]) -> int:
        """
        Populate Medication nodes from RxNorm data.
        
        Args:
            medications: List of medication dictionaries
            
        Returns:
            Number of medications created
        """
        query = """
        UNWIND $batch AS m
        MERGE (med:Medication:RxNormDrug {id: m.id})
        SET med.name = m.name,
            med.genericName = m.genericName,
            med.drugClass = m.drugClass,
            med.dosageForms = m.dosageForms,
            med.ontology = 'RxNorm',
            med.createdAt = datetime()
        RETURN count(*) as created
        """
        
        total = 0
        with self.driver.session() as session:
            for batch in self.batch_generator(medications):
                result = session.run(query, batch=batch)
                total += result.single()["created"]
        
        self.stats["nodes_created"] += total
        return total
    
    # =========================================================
    # POPULATE RELATIONSHIPS
    # =========================================================
    
    def populate_diagnoses(self, diagnoses: List[Dict]) -> int:
        """
        Create DIAGNOSED_WITH relationships.
        
        Args:
            diagnoses: List of diagnosis records
            
        Returns:
            Number of relationships created
        """
        query = """
        UNWIND $batch AS diag
        MATCH (p:Patient {id: diag.patientId})
        MATCH (d:Disease {id: diag.diseaseId})
        MERGE (p)-[r:DIAGNOSED_WITH]->(d)
        SET r.diagnosedDate = date(diag.diagnosedDate),
            r.diagnosedBy = diag.diagnosedBy,
            r.status = coalesce(diag.status, 'active'),
            r.severity = diag.severity,
            r.createdAt = datetime()
        RETURN count(*) as created
        """
        
        total = 0
        with self.driver.session() as session:
            for batch in self.batch_generator(diagnoses):
                result = session.run(query, batch=batch)
                total += result.single()["created"]
        
        self.stats["relationships_created"] += total
        return total
    
    def populate_prescriptions(self, prescriptions: List[Dict]) -> int:
        """
        Create TAKES relationships.
        
        Args:
            prescriptions: List of prescription records
            
        Returns:
            Number of relationships created
        """
        query = """
        UNWIND $batch AS rx
        MATCH (p:Patient {id: rx.patientId})
        MATCH (m:Medication {id: rx.medicationId})
        MERGE (p)-[r:TAKES]->(m)
        SET r.dosage = rx.dosage,
            r.frequency = rx.frequency,
            r.prescribedDate = date(rx.prescribedDate),
            r.prescribedBy = rx.prescribedBy,
            r.durationDays = rx.durationDays,
            r.refills = rx.refills,
            r.status = coalesce(rx.status, 'active'),
            r.createdAt = datetime()
        RETURN count(*) as created
        """
        
        total = 0
        with self.driver.session() as session:
            for batch in self.batch_generator(prescriptions):
                result = session.run(query, batch=batch)
                total += result.single()["created"]
        
        self.stats["relationships_created"] += total
        return total
    
    def populate_treatments(self, treatments: List[Dict]) -> int:
        """
        Create TREATS relationships between medications and diseases.
        
        Args:
            treatments: List of treatment records
            
        Returns:
            Number of relationships created
        """
        query = """
        UNWIND $batch AS t
        MATCH (m:Medication {id: t.medicationId})
        MATCH (d:Disease {id: t.diseaseId})
        MERGE (m)-[r:TREATS]->(d)
        SET r.effectiveness = t.effectiveness,
            r.firstLine = coalesce(t.firstLine, false),
            r.createdAt = datetime()
        RETURN count(*) as created
        """
        
        total = 0
        with self.driver.session() as session:
            for batch in self.batch_generator(treatments):
                result = session.run(query, batch=batch)
                total += result.single()["created"]
        
        self.stats["relationships_created"] += total
        return total
    
    # =========================================================
    # POPULATE FROM FILE
    # =========================================================
    
    def populate_from_json(self, filepath: str) -> Dict[str, int]:
        """
        Populate entire knowledge graph from JSON file.
        
        Expected JSON structure:
        {
            "patients": [...],
            "diseases": [...],
            "medications": [...],
            "diagnoses": [...],
            "prescriptions": [...],
            "treatments": [...]
        }
        """
        data = self.load_json_file(filepath)
        results = {}
        
        # Create nodes first
        if "patients" in data:
            results["patients"] = self.populate_patients(data["patients"])
            print(f"   Created {results['patients']} patients")
        
        if "diseases" in data:
            results["diseases"] = self.populate_diseases(data["diseases"])
            print(f"   Created {results['diseases']} diseases")
        
        if "medications" in data:
            results["medications"] = self.populate_medications(data["medications"])
            print(f"   Created {results['medications']} medications")
        
        # Then create relationships
        if "diagnoses" in data:
            results["diagnoses"] = self.populate_diagnoses(data["diagnoses"])
            print(f"   Created {results['diagnoses']} diagnoses")
        
        if "prescriptions" in data:
            results["prescriptions"] = self.populate_prescriptions(data["prescriptions"])
            print(f"   Created {results['prescriptions']} prescriptions")
        
        if "treatments" in data:
            results["treatments"] = self.populate_treatments(data["treatments"])
            print(f"   Created {results['treatments']} treatment relationships")
        
        return results
    
    # =========================================================
    # POPULATE DOCUMENTS FOR RAG
    # =========================================================
    
    def populate_documents(self, documents: List[Dict]) -> int:
        """
        Populate Document nodes for RAG.
        
        Args:
            documents: List of document dictionaries
            
        Returns:
            Number of documents created
        """
        query = """
        UNWIND $batch AS doc
        MERGE (d:Document {id: doc.id})
        SET d.title = doc.title,
            d.source = doc.source,
            d.content = doc.content,
            d.createdAt = datetime()
        RETURN count(*) as created
        """
        
        total = 0
        with self.driver.session() as session:
            for batch in self.batch_generator(documents):
                result = session.run(query, batch=batch)
                total += result.single()["created"]
        
        self.stats["nodes_created"] += total
        return total
    
    def create_entity_mentions(self) -> int:
        """
        Create MENTIONS relationships between documents/chunks and entities.
        Uses text matching to find entity mentions.
        """
        query = """
        MATCH (d:Document)
        MATCH (e)
        WHERE (e:Disease OR e:Medication OR e:Patient)
          AND e.name IS NOT NULL
          AND toLower(d.content) CONTAINS toLower(e.name)
        MERGE (d)-[r:MENTIONS]->(e)
        SET r.createdAt = datetime()
        RETURN count(*) as created
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            count = result.single()["created"]
        
        self.stats["relationships_created"] += count
        return count
    
    # =========================================================
    # VALIDATION
    # =========================================================
    
    def validate_data(self) -> Dict[str, Any]:
        """Validate the populated data."""
        queries = {
            "orphan_patients": """
                MATCH (p:Patient)
                WHERE NOT (p)-[:DIAGNOSED_WITH]->()
                  AND NOT (p)-[:TAKES]->()
                RETURN count(p) as count
            """,
            "orphan_medications": """
                MATCH (m:Medication)
                WHERE NOT ()-[:TAKES]->(m)
                  AND NOT (m)-[:TREATS]->()
                RETURN count(m) as count
            """,
            "missing_treatment_links": """
                MATCH (p:Patient)-[:TAKES]->(m:Medication)
                MATCH (p)-[:DIAGNOSED_WITH]->(d:Disease)
                WHERE NOT (m)-[:TREATS]->(d)
                RETURN count(DISTINCT p) as count
            """
        }
        
        results = {}
        with self.driver.session() as session:
            for name, query in queries.items():
                result = session.run(query)
                results[name] = result.single()["count"]
        
        return results
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get current database statistics."""
        query = """
        MATCH (n)
        WITH labels(n) AS lbls, count(*) AS cnt
        UNWIND lbls AS label
        WITH label, sum(cnt) AS count
        RETURN label, count
        ORDER BY count DESC
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            labels = {r["label"]: r["count"] for r in result}
            
            # Count relationships
            rel_query = """
            MATCH ()-[r]->()
            WITH type(r) AS relType, count(*) AS count
            RETURN relType, count
            ORDER BY count DESC
            """
            result = session.run(rel_query)
            relationships = {r["relType"]: r["count"] for r in result}
        
        return {
            "labels": labels,
            "relationships": relationships,
            "import_stats": self.stats
        }


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Populate Ontology - Demonstration")
    print("=" * 60)
    
    try:
        driver = get_driver()
        populator = OntologyPopulator(driver, batch_size=100)
        
        # Check for sample data file
        data_file = Path(__file__).parent.parent / "data" / "sample_healthcare_data.json"
        
        if data_file.exists():
            print("\n📌 Loading Sample Healthcare Data")
            print("-" * 40)
            
            results = populator.populate_from_json(str(data_file))
            
            print("\n📌 Creating Entity Mentions")
            print("-" * 40)
            mentions = populator.create_entity_mentions()
            print(f"   Created {mentions} MENTIONS relationships")
            
            print("\n📌 Validation Results")
            print("-" * 40)
            validation = populator.validate_data()
            for check, count in validation.items():
                status = "✅" if count == 0 else "⚠️"
                print(f"   {status} {check}: {count}")
            
            print("\n📌 Database Statistics")
            print("-" * 40)
            stats = populator.get_statistics()
            print("   Node Labels:")
            for label, count in stats["labels"].items():
                print(f"      {label}: {count}")
            print("   Relationships:")
            for rel_type, count in stats["relationships"].items():
                print(f"      {rel_type}: {count}")
        else:
            print(f"\n   Sample data file not found: {data_file}")
            print("   Creating inline sample data...")
            
            # Create sample data inline
            patients = [
                {"id": "P001", "name": "John Doe", "dateOfBirth": "1980-05-15", 
                 "gender": "Male", "bloodType": "O+", "allergies": ["Penicillin"]},
                {"id": "P002", "name": "Jane Smith", "dateOfBirth": "1992-08-22",
                 "gender": "Female", "bloodType": "A-", "allergies": []}
            ]
            
            diseases = [
                {"id": "E11", "name": "Type 2 Diabetes Mellitus", 
                 "category": "Endocrine", "severity": "moderate"},
                {"id": "I10", "name": "Essential Hypertension",
                 "category": "Cardiovascular", "severity": "mild"}
            ]
            
            medications = [
                {"id": "6809", "name": "Metformin", "genericName": "Metformin HCl",
                 "drugClass": "Biguanide", "dosageForms": ["Tablet"]},
                {"id": "29046", "name": "Lisinopril", "genericName": "Lisinopril",
                 "drugClass": "ACE Inhibitor", "dosageForms": ["Tablet"]}
            ]
            
            print(f"\n   Created {populator.populate_patients(patients)} patients")
            print(f"   Created {populator.populate_diseases(diseases)} diseases")
            print(f"   Created {populator.populate_medications(medications)} medications")
        
        driver.close()
        
    except Exception as e:
        print(f"   Error: {e}")
    
    print("\n" + "=" * 60)
    print("✅ Ontology population complete!")
    print("=" * 60)
